# -*- coding:utf-8 -*-
import os
import sys
import requests
import datetime
from Crypto.Cipher import AES
from binascii import b2a_hex, a2b_hex

# reload(sys)
# sys.setdefaultencoding('utf-8')


import importlib
importlib.reload(sys)

def download(url, savefile_path=''):
    # 创建download目录
    download_path = os.getcwd() + "\download"
    if not os.path.exists(download_path):
        os.mkdir(download_path)

    #新建日期文件夹，如果没有指定文件存储位置，就用年月日创建文件夹，否则用指定的文件夹
    if not savefile_path:
        savefile_path = os.path.join(download_path, datetime.datetime.now().strftime('%Y%m%d'))
    else:
        savefile_path = os.path.join(download_path, savefile_path)
    print("savefile_path = " + savefile_path)
    if not os.path.exists(savefile_path):
        os.mkdir(savefile_path)


    #如果不存在first.m3u8，就重新获取，并保存，否则证明之前已经获取过，直接读取就行了
    if not os.path.exists(os.path.join(download_path,'first.m3u8')):
        all_content = requests.get(url).text # 获取第一层M3U8文件内容
        with open( os.path.join(download_path,'first.m3u8'),'w' ) as f:
            f.write(all_content)
        print("save first.m3u8")
    else:
        print("first.m3u8 is exist,just read it")
        all_content = open( os.path.join(download_path,'first.m3u8'),'r' ).read()

    # 不是M3U8文件，直接报错退出
    if "#EXTM3U" not in all_content:
        raise BaseException("非M3U8的链接")
    #
    if "EXT-X-STREAM-INF" in all_content: # 第一层
        # 从第一层中读取所有的文件组成行列表 file_line
        print("we need to get the second line")
        file_line = all_content.split("\n")
        # 找到 .m3u8 ，拼出第二层的链接
        for line in file_line:
            if '.m3u8' in line:
                url = url.rsplit("/", 1)[0] + "/" + line # 拼出第二层m3u8的URL
                print("second line url = " + url)
                # 没存，就存一下，同first.m3u8的逻辑
                if not os.path.exists(os.path.join(download_path,'second.m3u8')):
                    all_content = requests.get(url).text
                    with open( os.path.join(download_path,'second.m3u8'),'w' ) as f:
                        f.write(all_content)
                    print("second.m3u8 has been saved")
                else:
                    all_content = open( os.path.join(download_path,'second.m3u8'),'r' ).read()
                break
    # 到此为止，all_content里面的内容已经更新成了包含所有文件名的最终文件
    #           url 更新成了下载最终m3u8的URL

    # 把里面的元素分割出来
    file_line = all_content.split("\n")

    unknow = True
    key = ""

    for line in file_line:
        # 先找到解密key
        if "#EXT-X-KEY" in line: # 找解密Key
            method_pos = line.find("METHOD")
            comma_pos = line.find(",")
            method = line[method_pos:comma_pos].split('=')[1]
            print("Decode Method：", method)

            uri_pos = line.find("URI")
            quotation_mark_pos = line.rfind('"')
            key_path = line[uri_pos:quotation_mark_pos].split('"')[1]

            key_url = url.rsplit("/", 1)[0] + "/" + key_path # 拼出key解密密钥URL
            res = requests.get(key_path)
            key = res.content

            with open( os.path.join(download_path,'key.key'),'wb' ) as f:
                f.write(key)

            cryptor = AES.new(key, AES.MODE_CBC, key)
            print( "key：" , key)
            break
            print("you cann't see me")
    # 再进行下载和解密
    for index, line in enumerate(file_line):    # 第二层
        if "EXTINF" in line: # 找ts地址并下载
            unknow = False
            pd_url = url.rsplit("/", 1)[0] + "/" + file_line[index + 1] # 拼出ts片段的URL

            c_fule_name = file_line[index + 1].rsplit("/", 1)[-1]
            # 如果文件已经存在，就下一个
            if os.path.exists( os.path.join(savefile_path, c_fule_name) ):
                print("file %s exist, next" % c_fule_name)
                continue
            # 网络不好的时候会失败，所以重复3次
            for i in range(3):
                # 获取视频内容
                print("get video " + datetime.datetime.now().strftime('%H:%M:%S'))
                res = requests.get(pd_url)
                try:
                    if len(key): # AES 解密，有key就是需要解密
                        with open(os.path.join(savefile_path, c_fule_name), 'ab') as f:
                            f.write(cryptor.decrypt(res.content))
                        print(c_fule_name + " success")
                        break
                    else:  # 没有key就不需要解密了，直接保存文件就行了
                        with open(os.path.join(savefile_path, c_fule_name), 'ab') as f:
                            f.write(res.content)
                            f.flush()
                        print(c_fule_name + " success")
                        break
                except:  # 网络不好导致下载失败，清除文件
                    print( str(i+1) + " download error, file: " + c_fule_name)
                    os.remove(os.path.join(savefile_path, c_fule_name))
            else:  # 下载失败，先跳过这个文件，后面重新执行的时候在重新下载
                print("download file has failed 3 times, jump it")
                #exit()

    if unknow:
        raise BaseException("未找到对应的下载链接")
    else:
        print( "下载完成")
        # merge_file(savefile_path)

def merge_file(path):
    os.chdir(path)
    cmd = "copy /b * new.ts"
    #os.system(cmd)
    #os.system('del /Q *.ts')
    #os.system('del /Q *.mp4')
    #os.rename("new.tmp", "new.mp4")

if __name__ == '__main__':
    # url = input("please input the url of index.m3u8 file:\n")
    # savefile_path = input("give me a dir name to save all the video:\n(or you can just press Enter and I will create one for you)\n")
    # url = r''
    # savefile_path = r''

    url = "https://video.hcyunshang.cn/20210215/cEkErZnB/1000kb/hls/index.m3u8";
    savefile_path = "./bug_file/"
    download(url,savefile_path)